# -*- coding: utf-8 -*-
# @Time    : 2021/8/2 11:43
# @Author  : ZSQ
# @Email   : zsq199170918@163.com
# @FileName: redis_case_filter.py
# @Software: PyCharm

from redis import StrictRedis
from pymongo import MongoClient
from HifoXf.utils.redis_pool import POOL


def get_filter_flag(db_name: str, col_name_li: list, set_name: str, mongo_h=None, mongo_p=None,
                    mongo_client=None) -> bool:
    """
    基于redis，根据案例id去重
    :param db_name: mongo库名
    :param col_name_li: mongo集合名，列表
    :param set_name: redis集合名
    :param mongo_h: mongo host
    :param mongo_p: mongo port
    :param mongo_client: mongo连接对象
    :return:
    """
    # mongo连接对象
    mongo_cli = MongoClient(host=mongo_h, port=mongo_p, ) if not mongo_client else mongo_client
    # redis连接对象
    redis_conn = StrictRedis(connection_pool=POOL)
    redis_pipeline = redis_conn.pipeline()
    try:
        # 遍历集合，mongo查询
        case_id_list = list()
        for col_name in col_name_li:
            mongo_col = mongo_cli[db_name][col_name]
            col_case_dict_list = list(mongo_col.find({}, {'projectId': 1, '_id': 0}))
            col_case_id_list = [case_dict['projectId'] for case_dict in col_case_dict_list if
                                case_dict.get('projectId', None)]
            case_id_list += col_case_id_list
        # 写入redis
        # todo 分布式时需优化,同时在爬取时只能删除一次
        redis_pipeline.multi()
        redis_pipeline.delete(set_name)
        # 如果数据库中有案例才添加
        if case_id_list:
            redis_pipeline.sadd(set_name, *case_id_list)
        redis_pipeline.execute()
        del case_id_list
    except Exception as e:
        flag = False
    else:
        flag = True
    finally:
        # 关闭连接对象
        mongo_cli.close()
        redis_conn.close()
    return flag


if __name__ == '__main__':
    get_filter_flag('newhours', ['beike_xf', ], 'beike_xf_case', mongo_h='127.0.0.1', mongo_p=27017, )
